# -*- coding: utf-8 -*-
"""
Created on Thu Nov 12 19:26:15 2020

@author: DELL
"""

import os
import torch
from torchvision import transforms
import cv2
import random
import numpy as np
import pandas as pd
# import SimpleITK as sitk
import matplotlib.pyplot as plt


def visualize(image, label):
    plt.subplot(1, 2, 1)
    plt.imshow(image.squeeze().cpu())
    plt.subplot(1, 2, 2)
    plt.imshow(label.squeeze().cpu())
    plt.show()


class DiagnoseDataset(torch.utils.data.Dataset):
    def __init__(self, csv_path, scan_dir, device='cpu'):
        if not os.path.exists(csv_path):
            raise RuntimeError(".csv path: "+ csv_path +" not found")
        
        if not os.path.isdir(scan_dir):
            raise RuntimeError("Scan directory: "+ scan_dir +" not found")
            
        self.csv_path = csv_path
        self.scan_dir = scan_dir
        
        self.image_filelist = []
        self.label_critical = []
        
        csv = pd.read_csv(csv_path)
        for row in csv.itertuples():
            patient_id = str(getattr(row, 'scan_id'))
            if os.path.exists(os.path.join(scan_dir, patient_id)):
                doclist = os.listdir(os.path.join(scan_dir, patient_id))
                for doc in doclist:
                    new_file = [os.path.join(scan_dir, patient_id, doc, x) for x in os.listdir(os.path.join(scan_dir, patient_id, doc))]
                    new_file = new_file[20: -15]
                    self.image_filelist += new_file
                    self.label_critical += len(new_file) * [getattr(row, 'Critical_illness')]
        
        self.device = device
        self.lung_mask_model = torch.load("lung_mask.pkl")
      
        self.length = len(self.image_filelist)
    
    
    def lung_mask(self, image):
        mask = self.lung_mask_model(torch.unsqueeze(image, dim=0).float().cuda()).to(self.device)
        mask = torch.argmax(mask.squeeze(), dim=0)
        mask = torch.unsqueeze(mask, dim=0)
        image = image * mask + 170 * (mask == 0)
        return image
    
        
    def __getitem__(self, index):
        if index >= self.length or index < 0:
            raise RuntimeError("index out of bounds.")
        
        image = cv2.imread(self.image_filelist[index])[..., 0]
        image = torch.from_numpy(image).to(self.device)
        image = torch.unsqueeze(image, dim=0)
        image = transforms.functional.resize(image, (512, 512))
        
        
        # if np.random.randint(0, 1) == 0:
        #     image = transforms.functional.hflip(image)
        
        # if np.random.randint(0, 1) == 0:
        #     image = transforms.functional.vflip(image)
            
        # scale = np.random.uniform(0.9, 1.1)
        # size = int(scale * 512)
        # image = transforms.functional.resize(image, size)
        # if size > 512:
        #     image = transforms.functional.center_crop(image, 512)
        # else:
        #     left = int((512 - size) / 2)
        #     right = (512 - size) - left
        #     image = transforms.functional.pad(image, [left, left, right, right], 170)
  
        # rot = random.randint(0, 3)
        # image = torch.rot90(image, rot, dims=(1, 2))


        # image = self.lung_mask(image)
        image = torch.squeeze(image)
        image = torch.stack([image, image, image], dim=0)
        image = image.float()
        
        label = self.label_critical[index]
        label = torch.tensor(label).to(self.device)
        
        return image, label
        
        
    def __len__(self):
        return self.length
        
        
        

class LabeledDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, label_dir, device='cpu'):
        if not os.path.isdir(image_dir):
            raise RuntimeError("Image directory: "+ image_dir +" not found")
        
        if not os.path.isdir(label_dir):
            raise RuntimeError("Label directory: "+ label_dir +" not found")
            
        self.image_dir = image_dir
        self.label_dir = label_dir
        
        self.image_filelist = []
        self.label_filelist = []
        label_doclist = os.listdir(label_dir)
        for doc in label_doclist:
            image_docpath = os.path.join(image_dir, doc)
            label_docpath = os.path.join(label_dir, doc)
            if os.path.isdir(label_docpath):
                self.image_filelist += [os.path.join(image_docpath, x).replace("png", "jpg") for x in os.listdir(label_docpath)]
                self.label_filelist += [os.path.join(label_docpath, x) for x in os.listdir(label_docpath)]
                
        self.length = len(self.image_filelist)
        self.device = device
        
        np.random.seed(random.randint(1, 10000))
        
        
    def __getitem__(self, index):
        if index >= self.length or index < 0:
            raise RuntimeError("index out of bounds.")
        
        image = cv2.imread(self.image_filelist[index])[..., 0]
        image = torch.from_numpy(image).to(self.device)
        image = torch.unsqueeze(image, dim=0)
        
        label = cv2.imread(self.label_filelist[index])[..., 0]
        label = torch.from_numpy(label).to(self.device)
        label = torch.unsqueeze(label, dim=0)
        
        
        if np.random.randint(0, 1) == 0:
            image = transforms.functional.hflip(image)
            label = transforms.functional.hflip(label)
        
        if np.random.randint(0, 1) == 0:
            image = transforms.functional.vflip(image)
            label = transforms.functional.vflip(label)
            
        scale = np.random.uniform(0.9, 1.1)
        size = int(scale * 512)
        image = transforms.functional.resize(image, size)
        label = transforms.functional.resize(label, size)
        if size > 512:
            image = transforms.functional.center_crop(image, 512)
            label = transforms.functional.center_crop(label, 512)
        else:
            left = int((512 - size) / 2)
            right = (512 - size) - left
            image = transforms.functional.pad(image, [left, left, right, right], 170)
            label = transforms.functional.pad(label, [left, left, right, right], 0)
  
        rot = random.randint(0, 3)
        image = torch.rot90(image, rot, dims=(1, 2))
        label = torch.rot90(label, rot, dims=(1, 2))


        image = image * (label > 0) + 170 * (label == 0)
        # label = label.squeeze()
        label = torch.clamp(label.int()-1, min=0, max=2)
        # label = (label > 0).long()
        
        return image.float(), label.long()
        
    def __len__(self):
        return self.length
    
    
    
class UnlabeledDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, label_dir, length, device='cpu'):
        if not os.path.isdir(image_dir):
            raise RuntimeError("Image directory: "+ image_dir +" not found")
        
        if not os.path.isdir(label_dir):
            raise RuntimeError("Label directory: "+ label_dir +" not found")
            
        self.image_dir = image_dir
        self.label_dir = label_dir
        
        self.image_filelist = []
        self.label_filelist = []
        label_doclist = os.listdir(label_dir)
        for doc in label_doclist:
            image_docpath = os.path.join(image_dir, doc)
            label_docpath = os.path.join(label_dir, doc)
            if os.path.isdir(image_docpath):
                self.image_filelist += [os.path.join(image_docpath, x) for x in os.listdir(image_docpath) 
                                        if not os.path.exists(os.path.join(label_docpath, x).replace(".jpg", ".png"))]
                self.label_filelist += [os.path.join(image_docpath, x) for x in os.listdir(image_docpath) 
                                        if os.path.exists(os.path.join(label_docpath, x).replace(".jpg", ".png"))]
        
        random.shuffle(self.image_filelist)
        
        if length > len(self.image_filelist):
            raise RuntimeError("Requested length out of bounds, don't have so much data")
        self.length = length
        self.device = device
        
        self.lung_mask_model = torch.load("lung_mask.pkl")
        
        np.random.seed(random.randint(1, 10000))
        
    
    def lung_mask(self, image):
        mask = self.lung_mask_model(torch.unsqueeze(image, dim=0).float().cuda()).to(self.device)
        mask = torch.argmax(mask.squeeze(), dim=0)
        mask = torch.unsqueeze(mask, dim=0)
        image = image * mask + 170 * (mask == 0)
        return image
        
        
    def __getitem__(self, index):
        index %= self.length
        # if index >= self.length or index < 0:
        #     raise RuntimeError("index out of bounds.")
        
        image = cv2.imread(self.image_filelist[index])[..., 0]
        image = torch.from_numpy(image).to(self.device)
        image = torch.unsqueeze(image, dim=0)
        
        
        if np.random.randint(0, 1) == 0:
            image = transforms.functional.hflip(image)
        
        if np.random.randint(0, 1) == 0:
            image = transforms.functional.vflip(image)
            
        scale = np.random.uniform(0.9, 1.1)
        size = int(scale * 512)
        image = transforms.functional.resize(image, size)
        if size > 512:
            image = transforms.functional.center_crop(image, 512)
        else:
            left = int((512 - size) / 2)
            right = (512 - size) - left
            image = transforms.functional.pad(image, [left, left, right, right], 170)
  
        rot = random.randint(0, 3)
        image = torch.rot90(image, rot, dims=(1, 2))


        image = self.lung_mask(image)
        
        return image
        
    def __len__(self):
        return self.length


def get_TCSM_dataloader(image_dir, label_dir, labeled_batch_size, valid_split=0.10, unlabeled_times=1, device='cpu'):
    labeled_dataset = LabeledDataset(image_dir, label_dir)
    
    valid_num = int(len(labeled_dataset) * valid_split)
    train_num = len(labeled_dataset) - valid_num
    trainset, validset = torch.utils.data.random_split(labeled_dataset, [train_num, valid_num])
    
    unlabeled_dataset = UnlabeledDataset(image_dir, label_dir, unlabeled_times*train_num)
    
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=labeled_batch_size, shuffle=True, drop_last=True)
    valid_loader = torch.utils.data.DataLoader(validset, batch_size=1, shuffle=True, drop_last=True)
    unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=labeled_batch_size*unlabeled_times, shuffle=True, drop_last=True)
    
    return train_loader, valid_loader, unlabeled_loader




if __name__ == "__main__":
    CSV_PATH = r'D:\data\metadata.csv'
    SCAN_DIR = r'D:\data\NCP'
    
    dataset = DiagnoseDataset(CSV_PATH, SCAN_DIR)
    
    
#     IMAGE_DIR = "D:\data\ct_lesion_seg\image"
#     MASK_DIR = "D:\data\ct_lesion_seg\mask"
#     # labeled_dataset = LabeledDataset(IMAGE_DIR, MASK_DIR)
#     # unlabeled_dataset = UnlabeledDataset(IMAGE_DIR, MASK_DIR, len(labeled_dataset))
#     # labeled_loader = torch.utils.data.DataLoader(labeled_dataset, batch_size=8, shuffle=True)
#     # unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=8, shuffle=True)
#     t, v, u = get_TCSM_dataloader(IMAGE_DIR, MASK_DIR, 5)
    